# -*- coding: utf-8 -*-
""""
selenium 辅助爬取
login方法打开登陆页面
run方法用线程挂起get_cookie函数循环判断cookie.txt文件直至爬虫退出
该程序在tianyancha.py开始时使用
"""
import os
import threading
from time import sleep
from selenium import webdriver
from pynput.keyboard import Key, Listener


class OpenDriver:
    """selenium登陆和获取cookie"""

    def __init__(self):
        self.driver = webdriver.Chrome(r'D:\软件存放\谷歌浏览器\Chrome\Application\chromedriver.exe')
        self.url = 'https://www.tianyancha.com/'
        self.verification_url = r'https://antirobot.tianyancha.com' \
                                r'/captcha/verify?return_url=https%3A%2F%2F' \
                                r'www.tianyancha.com%2Fcompany%2F25382265&rnd='
        self.first_handle = None

    def login(self):
        """登陆函数"""

        def on_press(key):
            """解除阻塞"""
            if str(key) == "'g'":                                           # 如果检测到键盘按下'G'
                print('关闭阻塞...')
                return False

        f = open('cookie.txt', mode='w')                                    # 保存新的cookie
        f.write('wait...')
        self.driver.get(self.url)                                           # 用selenium打开网页
        self.driver.maximize_window()                                       # 最大化窗口
        print("已阻塞..\n请登录..."
              "登完按'G'键解除阻塞...(网页加载完再'G')")
        self.first_handle = self.driver.current_window_handle               # 记录当前窗口句柄
        with Listener(on_press=on_press) as listener:                       # 开启键盘监控并阻塞
            listener.join()

        f = open('cookie.txt', mode='w')                                    # 保存新的cookie
        f.write(str(self.driver.get_cookies()))
        f.close()

    def get_cookie(self):
        """获取cookie"""
        while True:
            sleep(0.1)
            f = open(r'cookie.txt')
            data = f.read()
            if data == 'wait...':
                self.close_extra_windows()                                  # 关闭多余窗口
                self.driver.get(self.verification_url)                      # 用selenium打开验证码页面
                os.system('python ClickTheVerificationCode.py')             # 点击验证码
                html_str = str(self.driver.page_source)  # 获取源码
                if '我们只是确认一下你不是机器人' not in html_str:          # 判断是否点击成功
                    f = open('cookie.txt', mode='w')  # 保存新的cookie
                    f.write(str(self.driver.get_cookies()))
                    f.close()
                else:                                                       # 如果不成功重新点击
                    self.get_cookie()
            elif data == 'CloseSpider':
                break

    def close_extra_windows(self):
        """关闭多余窗口
        会莫名弹窗影响正常运行需要将对程序造成干扰的窗口关闭"""

        windows = self.driver.window_handles                                # 获取所有窗口句柄
        for handle in windows:                                              # 循环便利所有窗口
            if handle != self.first_handle:                                 # 如果不等于登陆时记录的窗口
                self.driver.switch_to.window(handle)                        # 切换到该窗口
                self.driver.close()                                         # 关闭它
            self.driver.switch_to.window(self.first_handle)                 # 切换回初始窗口

    def run(self):
        """挂起get_cookie函数"""
        t = threading.Thread(target=self.get_cookie)
        t.start()
